# -- coding: utf-8 --

"""
Author: Wuyang Yuan
Time: 2021-6-26 10:59:33
Description: 实现线性路网的列车运行图编制
路网：A->B->C->D->E->F...
列车：1, 2, 3, 4, 5, 6...
"""

from matplotlib import pyplot as plt
import numpy as np
from mip import Model, INTEGER, CONTINUOUS, BINARY, MINIMIZE, xsum
import random


# 铁路线路
class RailwayLine:
    def __init__(self, list):
        self.StationList = []
        self.SectionList = []
        self.PrevSection = {}
        self.NextSection = {}
        for name in list:
            if len(self.StationList) == 0:
                self.StationList.append(name)
            else:
                last = self.StationList[len(self.StationList) - 1]
                station = name
                section = Section(last, station)
                self.StationList.append(station)
                self.SectionList.append(section)
                self.PrevSection[station] = section
                self.NextSection[last] = section

    # 打印
    def Print(self):
        for station in self.StationList:
            if station == self.StationList[len(self.StationList) - 1]:
                print(station)
            else:
                print(station, '-->', end="")

    # 通过车站名称获取区间
    def GetSec(self, sta1, sta2):
        for sec in self.SectionList:
            if sec.start_station == sta1 and sec.end_station == sta2:
                return sec

    # 获取除去首末站的其它车站
    def GetInterStation(self):
        for station in self.StationList:
            if not (station == self.StationList[0] or station == self.StationList[len(self.StationList) - 1]):
                yield station


# 区间
class Section:
    def __init__(self, station1, station2):
        self.start_station = station1
        self.end_station = station2


# 时刻表
class Timetable:

    def __init__(self):
        self.Structure = {}
        self.TrainList = []

    # 根据列车数和路网生成空运行图
    def Generate(self, num, RailwayLine):
        for section in RailwayLine.SectionList:
            self.Structure[section] = []
        for i in range(num):
            t = Train(i)
            for sta in RailwayLine.StationList:
                if sta == RailwayLine.StationList[0] or sta == RailwayLine.StationList[
                    len(RailwayLine.StationList) - 1]:
                    t.AddStop(sta, 0, 0, True)
                else:
                    t.AddStop(sta, 0, 0, False)
            self.TrainList.append(t)

    # 通过Id获取列车
    def GetTrainById(self, trainId):
        return self.TrainList[trainId]

    # 根据时刻表获取结构
    def UpdateStructure(self, RailwayLine):
        for section in RailwayLine.SectionList:
            a = []
            for t in self.TrainList:
                a.append(t.GetStop(section.start_station).DepTime)
            self.Structure[section] = np.argsort(a).tolist()

    # 复制时刻表
    def copy(self):
        new_timetable = Timetable()
        for key in self.Structure:
            new_timetable.Structure[key] = self.Structure[key].copy()
        for train in self.TrainList:
            t = Train(train.Id)
            for stop in train.StopList:
                t.AddStop(stop.Station, stop.ArrTime, stop.DepTime, stop.Stop)
            new_timetable.TrainList.append(t)
        return new_timetable

    def print_structure(self, net):
        for s in net.SectionList:
            print(s.start_station, '->', s.end_station, self.Structure[s])


# 列车
class Train:
    def __init__(self, trainId):
        self.Id = trainId
        self.StopList = []

    # 增加停站
    def AddStop(self, station, arr_time, dep_time, isStop):
        self.StopList.append(Stop(station, arr_time, dep_time, isStop))

    # 获取停站
    def GetStop(self, name):
        for stop in self.StopList:
            if stop.Station == name:
                return stop

    # 获取区间运行时间
    def GetSecTimes(self):
        for index in range(len(self.StopList) - 1):
            yield [self.StopList[index].Station, self.StopList[index].DepTime,
                   self.StopList[index + 1].Station, self.StopList[index + 1].ArrTime]


# 列车停站
class Stop:
    def __init__(self, station, arr_time, dep_time, stop):
        self.Station = station
        self.ArrTime = arr_time
        self.DepTime = dep_time
        self.Stop = stop

    def IsPass(self):
        return self.ArrTime == self.DepTime


# 服务约束类
class ServiceConstraint:
    def __init__(self, num, stationStopLB, trainStopUB):
        self.NumOfTrains = num  # 开行列车总数
        self.StationStopLB = stationStopLB  # 最低停站频次
        self.TrainStopUB = trainStopUB  # 列车停站上限

    def SetStationLB(self, station, num):
        self.StationStopLB[station] = num

    def SetTrainStopUB(self, train, num):
        self.TrainStopUB[train] = num


# 参数类
class TimetableParameters:
    def __init__(self, rTime, iTime, aTime, sTime, bTime, hTime):
        self.RunningTime = rTime  # 区间运行时间
        self.MinIntervalTime = iTime  # 最短停站时间
        self.MaxIntervalTime = aTime
        self.StartAttachTime = sTime
        self.BarkeAttachTime = bTime
        self.Headway = hTime


# 绘图类
class TimetableView:
    def __init__(self, net):
        self.net = net
        self.station_location_dict = {}
        self.train_color_dict = {}
        last = 0
        for station in net.StationList:
            self.station_location_dict[station] = last + 1
            last += 1

        self.fig, self.axs = plt.subplots(1, 1, figsize=(6, 6))
        plt.xticks(np.arange(0, 480, 60))
        plt.yticks(np.arange(1, len(self.net.StationList) + 1), self.net.StationList)
        self.axs.set_xlim(0, 480)
        self.axs.set_xlabel('time   /min')
        self.axs.set_ylabel('station')
        self.axs.grid(True)

    # 绘制运行图
    def plot_timetable(self, timetable, **kargs):
        # Draw Train Lines
        for t in timetable.TrainList:
            if self.train_color_dict.__contains__(t.Id):
                color = self.train_color_dict[t.Id]
            else:
                color = "#"
                for i in range(6):  # 16进制表示颜色
                    color += random.choice('0123456789ABCDEF')
                self.train_color_dict[t.Id] = color
            for stop in t.GetSecTimes():
                self.axs.plot([stop[1], stop[3]], [self.station_location_dict[stop[0]],
                                                   self.station_location_dict[stop[2]]], color=color, label=str(t.Id))
        # 显示图形的图例
        # plt.legend()

        if kargs.__contains__('title'):
            plt.title(kargs['title'])

    # 绘制运行图
    def plot_cpm(self, timetable, path):
        # Draw CPM
        if path != '':
            for index in range(len(path) - 1):  # node = 0_A_0
                if path[index] == 'src' or path[index + 1] == 'end':
                    continue
                else:
                    str1 = path[index].split('_')
                    train1 = timetable.TrainList[int(str1[0])]
                    sta1 = str1[1]
                    da1 = str1[2]
                    x1 = train1.GetStop(sta1).DepTime if da1 == '1' else train1.GetStop(sta1).ArrTime

                    str2 = path[index + 1].split('_')
                    train2 = timetable.TrainList[int(str2[0])]
                    sta2 = str2[1]
                    da2 = str2[2]
                    x2 = train2.GetStop(sta2).DepTime if da2 == '1' else train2.GetStop(sta2).ArrTime

                    self.axs.plot([x1, x2], [self.station_location_dict[sta1], self.station_location_dict[sta2]],
                                  color='red', linewidth=5, linestyle=":")

    def clear(self):
        self.axs.lines.clear()


# 生成随机停站案例
def glb_random_generate_world(num_of_stations=10, num_of_trains=20):
    stations = []
    for i in range(num_of_stations):
        stations.append(chr(i + 65))
    net = RailwayLine(stations)
    StationStopLB = {}
    TrainStopUB = {}
    for sta in net.StationList:
        if sta == net.StationList[0] or sta == net.StationList[len(net.StationList) - 1]:
            StationStopLB[sta] = num_of_trains
        else:
            StationStopLB[sta] = random.randint(0, num_of_trains)
    for i in range(num_of_trains):
        TrainStopUB[i] = random.randint(2, num_of_stations)
    constraints = ServiceConstraint(num_of_trains, StationStopLB, TrainStopUB)
    RunningTime = {}
    Min_Interval_Time = {}
    Max_Interval_Time = {}
    for sec in net.SectionList:
        RunningTime[sec] = random.randint(10, 20)
    for sta in net.StationList:
        Min_Interval_Time[sta] = 2
        Max_Interval_Time[sta] = 30
    params = TimetableParameters(RunningTime, Min_Interval_Time, Max_Interval_Time, 2, 3, 5)
    return net, constraints, params


# 生成小案例
def glb_generate_test_world():
    net = RailwayLine(['A', 'B', 'C', 'D', 'E', 'F'])  # 线路上的停站
    StationStopLB = {'A': 5, 'B': 3, 'C': 4, 'D': 5, 'E': 4, 'F': 5}  # 每个车站的最小停站次数
    TrainStopUB = {0: 6, 1: 6, 2: 6, 3: 6, 4: 6}  # 列车最多停站次数
    constraints = ServiceConstraint(5, StationStopLB, TrainStopUB)
    RunningTime = {net.GetSec('A', 'B'): 10,  # 区间运行时间
                   net.GetSec('B', 'C'): 20,
                   net.GetSec('C', 'D'): 10,
                   net.GetSec('D', 'E'): 10,
                   net.GetSec('E', 'F'): 10}
    IntervalTime = {'B': 2,  # 车站间隔时间
                    'C': 2,
                    'D': 2,
                    'E': 2}
    params = TimetableParameters(RunningTime, IntervalTime, 2, 3,3, 5)
    return net, constraints, params


# 生成初始运行图
def glb_generate_basic_timetable(net, constraints, params):
    timetable = Timetable()
    timetable.Generate(constraints.NumOfTrains, net)
    glb_random_assign_stop(net, timetable, constraints)
    for sec in net.SectionList:
        for train in timetable.TrainList:
            timetable.Structure[sec].append(train.Id)
    return timetable


# 分配停站
def glb_random_assign_stop(network, timetable, constraints):
    station_stop_count = constraints.StationStopLB.copy()
    train_stop_count = constraints.TrainStopUB.copy()
    for sta in network.GetInterStation():
        count = constraints.StationStopLB[sta]  # 总停站
        alter = []
        for t in timetable.TrainList:
            if train_stop_count[t.Id] > 0:
                alter.append(t)
        for i in range(count):
            if len(alter) == 0:
                break
            train = random.choice(alter)
            train.GetStop(sta).Stop = True
            train_stop_count[train.Id] -= 1
            alter.remove(train)
        alter.clear()


# 根据区间顺序求解时刻 (PY-MIP)
def glb_solve_structure2(Timetable, TimetableParameters, RailwayLine):
    model = Model()
    variables = {}
    virtual_end = model.add_var(lb=0, ub=1440, var_type=CONTINUOUS)

    for t in Timetable.TrainList:
        # Generate Departure and Arrival time variable
        variables[t] = {(i, j): model.add_var(lb=0, ub=1440, var_type=CONTINUOUS, name="{0}_{1}_{2}".format(t.Id, i, j))
                        for j in range(2) for i in RailwayLine.StationList}
        # Add Running Time Constraints
        for sec in RailwayLine.SectionList:
            runningTime = TimetableParameters.RunningTime[sec]
            if t.GetStop(sec.start_station).Stop:
                runningTime += TimetableParameters.StartAttachTime
            if t.GetStop(sec.end_station).Stop:
                runningTime += TimetableParameters.BarkeAttachTime
            model += variables[t][sec.end_station, 0] - variables[t][sec.start_station, 1] == runningTime
        # Add Dwelling Time Constraints
        for sta in RailwayLine.StationList:
            if sta == RailwayLine.StationList[0]:
                model += variables[t][sta, 0] == variables[t][sta, 1]
            elif sta == RailwayLine.StationList[len(RailwayLine.StationList) - 1]:
                model += variables[t][sta, 0] == variables[t][sta, 1]
                model += variables[t][sta, 1] <= virtual_end
            elif t.GetStop(sta).Stop:
                model += variables[t][sta, 1] - variables[t][sta, 0] >= TimetableParameters.MinIntervalTime[sta]
                # model.add_constraint(variables[t][sta, 1] - variables[t][sta, 0] <= TimetableParameters.MaxIntervalTime[sta])

            else:
                model += variables[t][sta, 0] == variables[t][sta, 1]
    # Add Headway Constraints
    for sec in RailwayLine.SectionList:
        for i in range(len(Timetable.Structure[sec]) - 1):
            first_train = Timetable.TrainList[Timetable.Structure[sec][i]]
            second_train = Timetable.TrainList[Timetable.Structure[sec][i + 1]]
            model += variables[second_train][sec.start_station, 1] - variables[first_train][
                sec.start_station, 1] >= TimetableParameters.Headway  # headway
            model += variables[second_train][sec.end_station, 0] - variables[first_train][
                sec.end_station, 0] >= TimetableParameters.Headway  # headway

    # Add Obj Function
    model.objective = virtual_end
    model.optimize()

    if model.num_solutions:  # Get Value from result
        for t in Timetable.TrainList:
            for sta in RailwayLine.StationList:
                t.GetStop(sta).ArrTime = variables[t][sta, 0].x
                t.GetStop(sta).DepTime = variables[t][sta, 1].x
        return True, format(model.objective_value)
    else:
        return False, -1


def glb_solve_mip(timetable, net, constraints, params):
    model = Model()
    time_var = {}
    stop_var = {}
    order_var = {}
    # 定义变量
    # t, y
    for t in timetable.TrainList:
        # Generate Departure and Arrival time variable i=车站, j=到发 t=列车
        time_var[t] = {(i, j): model.add_var(lb=0, ub=1440, var_type=CONTINUOUS, name="{0}_{1}_{2}".format(t.Id, i, j))
                       for j in range(2) for i in net.StationList}
        stop_var[t] = {i: model.add_var(var_type=BINARY, name="{0}_{1}".format(t.Id, i))
                       for i in net.StationList}  # y=1停站 y=0不停

    # o
    for s in net.SectionList:
        order_var[s] = {}
        for t in timetable.TrainList:
            order_var[s][t] = {}
            for k in timetable.TrainList:
                if t.Id < k.Id:
                    order_var[s][t][k] = model.add_var(var_type=BINARY, name="{0}_{1}_{2}".format(s, t.Id, k.Id))

    # 辅助变量
    virtual_end = model.add_var(lb=0, ub=1440, var_type=CONTINUOUS)

    # 定义约束
    for t in timetable.TrainList:
        # 区间运行时间约束
        for sec in net.SectionList:
            model += time_var[t][sec.end_station, 0] - time_var[t][sec.start_station, 1] == params.RunningTime[sec] + \
                     params.StartAttachTime * stop_var[t][sec.start_station] + stop_var[t][sec.end_station]
        # 停站时间约束
        for sta in net.StationList:
            if sta == net.StationList[0]:  # 始发站
                model += stop_var[t][sta] == 1
                model += time_var[t][sta, 0] == time_var[t][sta, 1]
            elif sta == net.StationList[len(net.StationList) - 1]:  # 终到站
                model += stop_var[t][sta] == 1
                model += time_var[t][sta, 0] == time_var[t][sta, 1]
                model += time_var[t][sta, 1] <= virtual_end
            else:
                model += time_var[t][sta, 1] - time_var[t][sta, 0] >= params.MinIntervalTime[sta] * stop_var[t][sta]

        # 列车停站上限约束
        model += xsum(stop_var[t][sta] for sta in net.StationList) <= constraints.TrainStopUB[t.Id]

    for s in net.StationList:
        if s != net.StationList[0] or s != net.StationList[len(net.StationList) - 1]:
            model += xsum(stop_var[t][s] for t in timetable.TrainList) >= constraints.StationStopLB[s]  # 每个车站的停站下限

    # 车站间隔时间约束 - 只考虑到到和发发
    for sec in net.SectionList:
        for t in timetable.TrainList:
            for k in timetable.TrainList:
                if t.Id < k.Id:  # 确保只遍历一遍
                    model += time_var[t][sec.start_station, 1] - time_var[k][sec.start_station, 1] + (1 - order_var[sec][t][k]) * 1440 >= params.Headway
                    model += time_var[k][sec.start_station, 1] - time_var[t][sec.start_station, 1] + order_var[sec][t][k] * 1440 >= params.Headway  # 发发 headway 引入M = 1440
                    model += time_var[t][sec.end_station, 0] - time_var[k][sec.end_station, 0] + (1 - order_var[sec][t][k]) * 1440 >= params.Headway
                    model += time_var[k][sec.end_station, 0] - time_var[t][sec.end_station, 0] + order_var[sec][t][k] * 1440 >= params.Headway  # 到到 headway 引入M = 1440

    # 目标函数
    model.objective = virtual_end

    # 求解
    model.write("model.lp")
    model.max_mip_gap = 0.01
    model.optimize(max_seconds=60)

    makespan = 0
    # 输出
    if model.num_solutions:
        for t in timetable.TrainList:
            for sta in net.StationList:
                t.GetStop(sta).ArrTime = time_var[t][sta, 0].x
                t.GetStop(sta).DepTime = time_var[t][sta, 1].x
        makespan = format(model.objective_value)

    return 0, makespan


if __name__ == '__main__':
    net, constraints, params = glb_generate_test_world()  # 生成问题
    # net, constraints, params = glb_random_generate_world(4, 12)  # 生成问题
    timetable = glb_generate_basic_timetable(net, constraints, params)
    isSolved, makespan = glb_solve_mip(timetable, net, constraints, params)
    view = TimetableView(net)
    plt.ion()
    view.plot_timetable(timetable, title='Origin Timetable, Makespan=' + str(makespan))
    plt.savefig('origin.png', dpi=300)
    plt.show(block=True)
